#include <iostream>
#include<vector>
#include<algorithm>
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>

#include <fstream>

// 返回结果的结构体
struct Detection {
    cv::Rect box;
    float conf;
    size_t clas_id;
};

class YoloDetect {
public:
//    explicit YoloDetect(std::nullptr_t) {};
    // 构造函数初始化模型
    YoloDetect(const std::string &onnx_path, const float &conf_threshold);

    // 侦测，返回结构体[box,conf,clas_id]
    std::vector<Detection> detect(cv::Mat &draw_img);

private:
    Ort::Env env{nullptr};
    Ort::SessionOptions sessionOptions;
    // 初始化onnx模型
    Ort::Session session{nullptr};

    // 缩放图像加黑边
    void ScaleResize(const cv::Mat &image, cv::Mat &resize_image, std::map<std::string, float> &m_resize);

    // BGR --> RGB 标准化 除255
    void ImgNormalize(const cv::Mat &resize_image, cv::Mat &float_image);

    // BGR --> RGB 标准化 减均值除方差
    void ImgNormalize_(const cv::Mat &resize_image, cv::Mat &float_image);

    // 输出处理：置信度阈值筛选，坐标反算，nms
    std::vector<Detection> ScaleNms(const std::vector<Ort::Value> &outputTensors, const cv::Size &realimg_shape,
                                    const std::map<std::string, float> &m_resize);

    // 置信度阈值
    float confthreshold;

    //获取输入name;
    std::vector<const char *> input_names;
    //获取输出name;
    std::vector<const char *> output_names;

    // 网络需要的输入形状(某个轴可能是动态)，int64_t固定占用8个字节，不受编译器和操作系统影响，为了考虑可移植性
    std::vector<int64_t> inputtensor_shape;
    // 需要输入图像的 w,h  类型为float
    cv::Size2f input_size2f;
};

//==========================utils===========================
// 输出结果后画框
std::vector<cv::Scalar> SetColor(const size_t &num_classes);

void DrawImg(cv::Mat &img, const std::vector<Detection> &result, const std::vector<std::string> &class_names);


// =======================导出接口============================
struct Detection_ {
    float conf;  // 类别置信度
    size_t clas_id;  // 类别索引
    int *box_array;  // 存放4个坐标数组的指针 [x,y,w,h]
};

Detection_ Box2Array(cv::Rect box, const size_t &clas_id, const float &conf);

cv::Mat data2img(int height, int width, const unsigned char *img_data);

// 需要导出的函数
extern "C"
{
__declspec(dllexport) YoloDetect *InitDetect(const char *engine_file, const float &conf_threshold);
__declspec(dllexport) Detection_ *
OutResult(YoloDetect *ud, int &height, int &width, const unsigned char *img_data, size_t &result_size);
__declspec(dllexport) void FreeResult(Detection_ *result, size_t &result_size);
}
